package main

import (
	"fmt"
	"github.com/prometheus/prometheus/model/labels"
	"log"

	//"github.com/prometheus/prometheus/pkg/labels"
	"github.com/prometheus/prometheus/promql/parser"
)

func main() {
	expr := `sum by(cluster,node,role,instance)(100-(node_memory_MemAvailable_bytes/node_memory_MemTotal_bytes*100) * on(instance) group_left(role,node)(label_replace(kube_node_info * on(node) group_left(role) kube_node_role{role="master"},"instance","$1:9101", "internal_ip", "(.*)"))) >90`

	parsedExpr, err := parser.ParseExpr(expr)
	if err != nil {
		log.Fatalf("Failed to parse PromQL expression: %v", err)
	}
	fmt.Printf("Parsed expression: %v\n", parsedExpr)
	// 要解析的 PromQL 表达式

	// 遍历并打印解析后的 AST 节点
	traverseNode(parsedExpr, 0)
	// 注入标签
	injectLabels(parsedExpr, labels.MatchEqual, "cluster", "clusterName")
	fmt.Printf("Parsed expression: %v\n", parsedExpr)

}

func injectLabels(expr parser.Expr, match labels.MatchType, name, value string) {
	switch e := expr.(type) {
	case *parser.AggregateExpr:
		injectLabels(e.Expr, match, name, value)
	case *parser.Call:
		for _, v := range e.Args {
			injectLabels(v, match, name, value)
		}
	case *parser.ParenExpr:
		injectLabels(e.Expr, match, name, value)
	case *parser.UnaryExpr:
		injectLabels(e.Expr, match, name, value)
	case *parser.BinaryExpr:
		injectLabels(e.LHS, match, name, value)
		injectLabels(e.RHS, match, name, value)
	case *parser.VectorSelector:
		l := genMetricLabel(match, name, value)
		e.LabelMatchers = append(e.LabelMatchers, l)
		return
	case *parser.MatrixSelector:
		injectLabels(e.VectorSelector, match, name, value)
	case *parser.SubqueryExpr:
		injectLabels(e.Expr, match, name, value)
	case *parser.NumberLiteral, *parser.StringLiteral:
		return
	default:
		fmt.Printf("Unknown node type: %T\n", e)
	}
	return
}

func genMetricLabel(match labels.MatchType, name, value string) *labels.Matcher {
	m, err := labels.NewMatcher(match, name, value)
	if nil != err {
		return nil
	}

	return m
}

// 递归遍历 AST 节点并打印信息
func traverseNode(node parser.Node, indent int) {
	for i := 0; i < indent; i++ {
		fmt.Print("  ")
	}
	fmt.Printf("%T\n", node)

	switch n := node.(type) {
	case *parser.VectorSelector:
		for i := 0; i < indent+1; i++ {
			fmt.Print("  ")
		}
		fmt.Printf("Metric: %s\n", n.Name)
		for _, matcher := range n.LabelMatchers {
			for i := 0; i < indent+2; i++ {
				fmt.Print("  ")
			}
			fmt.Printf("LabelMatcher: %s %s %s\n", matcher.Name, matcher.Type, matcher.Value)
		}
	case *parser.MatrixSelector:
		for i := 0; i < indent+1; i++ {
			fmt.Print("  ")
		}
		fmt.Printf("Inner VectorSelector: %v\n", n.VectorSelector)
	}

	// 递归遍历子节点
	for _, child := range parser.Children(node) {
		traverseNode(child, indent+1)
	}
}
